
#include<vector>
#include<string>
#include<algorithm>
#include<cmath>
#include<complex>

#include<formic/command/command.h>
#include<formic/wfn/command.h>
#include<formic/mpi/interface.h>
#include<formic/random/random.h>
#include<formic/all_data.h>
#include<formic/exception.h>
#include<formic/numeric/numeric.h>
#include<formic/vector/matrix.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnIntData" command, which reads integer data from the
///          archive and adds it to the wavefunction
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_add_wfn_integer_data(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real")
    data.wd.add_wfn_int_data(arch);
  else if (data.userinp.get<std::string>("scalar_type") == "complex")
    data.wc.add_wfn_int_data(arch);
  else 
    throw formic::Exception("Unknown scalar type \"%s\" in formic_add_wfn_integer_data") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnRealData" command, which reads real data from the
///          archive and adds it to the wavefunction
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_add_wfn_real_data(formic::Archive & arch, formic::AllData & data) {

    data.wd.add_wfn_float_data(arch);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnRealDataUnif" command, which adds a real wfn data object
///          with uniform entries to the wavefunction
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_WfnRealDataUnif_command_func(formic::Archive & arch, formic::AllData & data) {

  // read name
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read name in formic_WfnRealDataUnif_command_func.");

  // read length
  int length;
  formic::mpi::read_and_bcast(arch, length, "Failed to read length in formic_WfnRealDataUnif_command_func.");

  // read value
  double val;
  formic::mpi::read_and_bcast(arch, val, "Failed to read value in formic_WfnRealDataUnif_command_func.");

  // read fixed/variable
  bool fixed;
  formic::mpi::read_and_bcast(arch, fixed, "Failed to read fixed in formic_WfnRealDataUnif_command_func.");

  // create the data object
  boost::shared_ptr<formic::WfnData<double> > d( new formic::WfnData<double>(name, fixed, length, val) );

  // add the data to the wavefunction
  data.wd.add_wfn_data(d);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnRealDataRand" command, which adds a real wfn data object
///          with random entries to the wavefunction
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_WfnRealDataRand_command_func(formic::Archive & arch, formic::AllData & data) {

  // read name
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read name in formic_WfnRealDataRand_command_func.");

  // read length
  int length;
  formic::mpi::read_and_bcast(arch, length, "Failed to read length in formic_WfnRealDataRand_command_func.");

  // read fixed/variable
  bool fixed;
  formic::mpi::read_and_bcast(arch, fixed, "Failed to read fixed in formic_WfnRealDataRand_command_func.");

  // create the data object
  boost::shared_ptr<formic::WfnData<double> > d( new formic::WfnData<double>(name, fixed, length, 0.0) );

  // randomize the entries
  std::vector<double> & vec = d->data_vec();
  for (int i = 0; i < vec.size(); i++)
    vec[i] = 2.0 * formic::random_number<double>() - 1.0;

  // broadcast the entries to all processes
  formic::mpi::bcast(&vec.at(0), vec.size());

  // add the data to the wavefunction
  data.wd.add_wfn_data(d);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnComplexData" command, which reads complex data from the
///          archive and adds it to the wavefunction
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_add_wfn_complex_data(formic::Archive & arch, formic::AllData & data) {

  data.wc.add_wfn_float_data(arch);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Adds an expansio of a smaller matrix to a larger matrix.
///
/// \param[in]       userinp  the user's input options
/// \param[in,out]   wfn      the wavefunction object holding the matrices
/// \param[in,out]   arch     archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
static void formic_ExpandMatrix_command_function_detail(const formic::InputBase & userinp,
                                                        formic::Wavefunction<S> & wfn,
                                                        formic::Archive & arch) {

  // variables to describe the two sets of jastrow factors
  int lg_dim, sm_dim; // number of orbitals
  std::string lg_name, sm_name;

  //// read in the indexing scheme ( 'c' for c, 'f' for fortran )
  //char indexing_scheme;
  //formic::mpi::read_and_bcast(arch, indexing_scheme, "Failed to read indexing_scheme in formic_ExpandMatrix_command_function_detail.");

  char indexing_scheme = 'f';

  // read in the dimension and name of the smaller matrix
  formic::mpi::read_and_bcast(arch, sm_dim, "Failed to read sm_dim in formic_ExpandMatrix_command_function_detail.");
  formic::mpi::read_and_bcast(arch, sm_name, "Failed to read sm_name in formic_ExpandMatrix_command_function_detail.");

  // read in the dimension and name of the larger matrix
  formic::mpi::read_and_bcast(arch, lg_dim, "Failed to read lg_dim in formic_ExpandMatrix_command_function_detail.");
  formic::mpi::read_and_bcast(arch, lg_name, "Failed to read lg_name in formic_ExpandMatrix_command_function_detail.");

  // read in the index mapping
  std::vector<int> imap;
  formic::mpi::read_and_bcast(arch, imap, "Failed to read index map in formic_ExpandMatrix_command_function_detail.");

  // get the small matrix
  const std::vector<S> & sm_mat = wfn.get_wf(sm_name)->data_vec();

  // get the large matrix
  std::vector<S> & lg_mat = wfn.get_wf(lg_name)->data_vec();

  // make sure large matrix is at least as large as small matrix
  if ( lg_dim < sm_dim ) throw formic::Exception("lg_dim must be >= sm_dim in formic_ExpandMatrix_command_function_detail");

  // check sizes
  if ( sm_mat.size() != sm_dim*sm_dim ) throw formic::Exception("sm_mat is the wrong size in formic_ExpandMatrix_command_function_detail");
  if ( lg_mat.size() != lg_dim*lg_dim ) throw formic::Exception("lg_mat is the wrong size in formic_ExpandMatrix_command_function_detail");

  // check index map size
  if ( imap.size() != sm_dim )
    throw formic::Exception("size of index map is %i, but should be %i in formic_ExpandMatrix_command_function_detail") % imap.size() % sm_dim;

  // check index map contents
  for (std::vector<int>::const_iterator it = imap.begin(); it != imap.end(); it++) {
    if ( *it < 0 || *it >= lg_dim )
      throw formic::Exception("out of range value \"%i\" in index map in formic_ExpandMatrix_command_function_detail") % *it;
    for (std::vector<int>::const_iterator jt = imap.begin(); jt != it; jt++)
      if ( *jt == *it )
        throw formic::Exception("repeated index \"%i\" in index map in formic_ExpandMatrix_command_function_detail") % *jt;
  }

  // (fortran indexing) add the expansion of the smaller matrix to the larger matrix
  if ( indexing_scheme == 'f' || indexing_scheme == 'F' ) {

    for (int j = 0; j < sm_dim; j++)
    for (int i = 0; i < sm_dim; i++)
      lg_mat[ imap[i] + imap[j] * lg_dim ] += sm_mat[ i + j * sm_dim ];

  // (c indexing) add the expansion of the smaller matrix to the larger matrix
  } else if ( indexing_scheme == 'c' || indexing_scheme == 'C' ) {

    for (int i = 0; i < sm_dim; i++)
    for (int j = 0; j < sm_dim; j++)
      lg_mat[ lg_dim * imap[i] + imap[j] ] += sm_mat[ sm_dim * i + j ];

  } else
    throw formic::Exception("unknown indexing_scheme \"%s\" in formic_ExpandMatrix_command_function_detail") % indexing_scheme;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "ExpandMatrix" command, which adds the expantion of a smaller
///          matrix to a larger matrix among the wavefunction data objects.
///
///          The first two arguments are the dimension and name of the smaller matrix.
///
///          The next two arguments are the dimension and name of the larger matrix.
///
///          The last argument is a vector giving the position of each smaller-matrix index
///          in the ordering of the larger matrix's indices.
///
///          For example, expanding the 4x4 matrix sm_mat into the lower right hand corner of
///          the 10x10 matrix lg_mat would use the command:
///
///            ExpandMatrix  4 sm_mat  10 lg_mat  [ 6 7 8 9 ]
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_ExpandMatrix_command_function(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic_ExpandMatrix_command_function_detail(data.userinp, data.wd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic_ExpandMatrix_command_function_detail(data.userinp, data.wc, arch);
  } else 
    throw formic::Exception("Unknown scalar type \"%s\" when processing ExpandMatrix command.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Adds random noise to a data object
///
/// \param[in]       userinp  the user's input options
/// \param[in,out]   wfn      the wavefunction object holding the matrices
/// \param[in,out]   arch     archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
static void formic_WfnDataNoise_command_function_detail(const formic::InputBase & userinp,
                                                        formic::Wavefunction<S> & wfn,
                                                        formic::Archive & arch) {

  // read in the magnitude of the noise
  double mag;
  formic::mpi::read_and_bcast(arch, mag, "Failed to read noise magnitude in formic_WfnDataNoise_command_function_detail.");

  // read in the name of the data object
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read data object name in formic_WfnDataNoise_command_function_detail.");

  // get a reference to the data vector
  std::vector<S> & vec = wfn.get_wf(name)->data_vec();

  // add the noise
  for (int i = 0; i < vec.size(); i++)
    vec[i] += mag * 2.0 * (              formic::unity(S()) * ( formic::random_number<double>() - 0.5 )
                             + formic::imaginary_unity(S()) * ( formic::random_number<double>() - 0.5 ) );

  // broadcast the data vector to all processes
  formic::mpi::bcast(vec);

  // empty any exact sample data in the wave function that may now have been invalidated
  wfn.clear_exact_sample_vecs();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnDataNoise" command, which adds random noise to a
///          wave function data object.
///
///          The first argument is the magnitude of the noise to add
///
///          The second argument is the name of the data object to add it to
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_WfnDataNoise_command_function(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic_WfnDataNoise_command_function_detail(data.userinp, data.wd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic_WfnDataNoise_command_function_detail(data.userinp, data.wc, arch);
  } else 
    throw formic::Exception("Unknown scalar type \"%s\" when processing WfnDataNoise command.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Scale a data object by a constant
///
/// \param[in]       userinp  the user's input options
/// \param[in,out]   wfn      the wavefunction object holding the matrices
/// \param[in,out]   arch     archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
static void formic_WfnDataScale_command_function_detail(const formic::InputBase & userinp,
                                                        formic::Wavefunction<S> & wfn,
                                                        formic::Archive & arch) {

  // read in the scaling constant
  S mag;
  formic::mpi::read_and_bcast(arch, mag, "Failed to read scaling constant in formic_WfnDataScale_command_function_detail.");

  // read in the name of the data object
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read data object name in formic_WfnDataScale_command_function_detail.");

  // get a reference to the data vector
  std::vector<S> & vec = wfn.get_wf(name)->data_vec();

  // scale the data
  for (int i = 0; i < vec.size(); i++)
    vec[i] *= mag;

  // broadcast the data vector to all processes
  formic::mpi::bcast(vec);

  // empty any exact sample data in the wave function that may now have been invalidated
  wfn.clear_exact_sample_vecs();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnDataScale" command, which scales the specified data
///          by the given constant.
///
///          The first argument is the magnitude of the scaling constant
///
///          The second argument is the name of the data object to add it to
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_WfnDataScale_command_function(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic_WfnDataScale_command_function_detail(data.userinp, data.wd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic_WfnDataScale_command_function_detail(data.userinp, data.wc, arch);
  } else 
    throw formic::Exception("Unknown scalar type \"%s\" when processing WfnDataScale command.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   symmetrizes the matrix stored in a wave function data object
///
/// \param[in]       userinp  the user's input options
/// \param[in,out]   wfn      the wavefunction object holding the matrices
/// \param[in,out]   arch     archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
static void formic_WfnDataSymmetrize_command_function_detail(const formic::InputBase & userinp,
                                                             formic::Wavefunction<S> & wfn,
                                                             formic::Archive & arch) {

  // read in the name of the data object
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read data object name in formic_WfnDataSymmetrize_command_function_detail.");

  // get a reference to the data vector
  std::vector<S> & vec = wfn.get_wf(name)->data_vec();

  // if the data vector is empty or length one then there is nothing to do
  if ( vec.size() < 2 )
    return;

  // get the matrix dimension
  int n = 0;
  for ( ; n < vec.size(); n++) {
    if ( n * n == vec.size() )
      break;
    if ( n * n > vec.size() )
      throw formic::Exception("The data in \"%s\" cannot be symmetrized because it is not a square matrix.") % name;
  }

  if ( formic::mpi::rank() == 0 ) {

    // construct a matrix around the data
    formic::Matrix<S,int> mat(n, n, &vec.at(0));

    // symmetrize the matrix
    mat += mat.t();
    mat *= 0.5 * formic::unity(S());

  }

  // broadcast the data vector to all processes
  formic::mpi::bcast(vec);

  // empty any exact sample data in the wave function that may now have been invalidated
  wfn.clear_exact_sample_vecs();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "WfnDataSymmetrize" command, which symmetrizes the matrix
///          stored in the supplied wave function data object.
///
///          The only argument is the name of the data object matrix to symmetrize
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_WfnDataSymmetrize_command_function(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic_WfnDataSymmetrize_command_function_detail(data.userinp, data.wd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic_WfnDataSymmetrize_command_function_detail(data.userinp, data.wc, arch);
  } else 
    throw formic::Exception("Unknown scalar type \"%s\" when processing WfnDataSymmetrize command.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "PrintVar" command, which prints a coefficient object's
///          variables.
///
///          The only argument is the name of the coefficient object.
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_PrintVar_command_function(formic::Archive & arch, formic::AllData & data) {

  // read the name of the coefficient object to use
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read coefficient object name in formic_PrintVar_command_function.");

  if ( formic::mpi::rank() == 0 ) {
    if (data.userinp.get<std::string>("scalar_type") == "real") {
      data.wd.get_co(co_name)->print_variables();
    } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
      data.wc.get_co(co_name)->print_variables();
    } else 
      throw formic::Exception("Unknown scalar type \"%s\" when processing PrintVar command.") % data.userinp.get<std::string>("scalar_type");
  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Add wavefunction related commands to the formic command system
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::init_wfn_commands() {

  // add command to read integer data
  formic::add_command("WfnIntData", &formic_add_wfn_integer_data);

  // add command to read real data
  formic::add_command("WfnRealData", &formic_add_wfn_real_data);

  // add command to read complex data
  formic::add_command("WfnComplexData", &formic_add_wfn_complex_data);

  formic::add_command("WfnRealDataUnif", &formic_WfnRealDataUnif_command_func);
  formic::add_command("WfnRealDataRand", &formic_WfnRealDataRand_command_func);

  formic::add_command("ExpandMatrix", &formic_ExpandMatrix_command_function);

  formic::add_command("WfnDataNoise", &formic_WfnDataNoise_command_function);

  formic::add_command("WfnDataSymmetrize", &formic_WfnDataSymmetrize_command_function);

  formic::add_command("WfnDataScale", &formic_WfnDataScale_command_function);

  formic::add_command("PrintVar", &formic_PrintVar_command_function);

}
